package com.wutong.resonance.model.util;

import lombok.extern.slf4j.Slf4j;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;

/**
 * 音频数据工具类，为音频数据转换提供便利
 *
 * @author wutong
 */
@Slf4j
public class AudioDataUtil {

    /**
     * float数组转byte数组
     *
     * @param data      float数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return byte数组
     */
    public static byte[] floatDataToByteArray(float[] data, ByteOrder byteOrder) {
        int length = data.length;
        if (length % 2 != 0) {
            length++;
        }
        ByteBuffer outBuffer = ByteBuffer.allocate(length * 2);
        outBuffer.order(byteOrder);

        for (int i = 0; i < data.length; i++) {
            if (data[i] > 1f) {
                data[i] = 1f;
            }
            if (data[i] < -1) {
                data[i] = -1;
            }
            outBuffer.putShort((short) (data[i] * 32767f));
        }
        return outBuffer.array();
    }

    /**
     * double数组转byte数组
     *
     * @param data      double数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return byte数组
     */
    public static byte[] doubleDataToByteArray(double[] data, ByteOrder byteOrder) {
        int length = data.length;
        if (length % 2 != 0) {
            length++;
        }
        ByteBuffer buf = ByteBuffer.allocate(length * 2);
        buf.order(byteOrder);

        for (int i = 0; i < length; i++) {
            if (data[i] > 1f) {
                data[i] = 1f;
            }
            if (data[i] < -1f) {
                data[i] = -1f;
            }
            buf.putShort((short) (data[i] * 32767f));
        }
        return buf.array();
    }

    /**
     * 二维float数组转byte数组
     *
     * @param matrix    二维float数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return byte数组
     */
    public static byte[] floatMatrixToByteArray(float[][] matrix, ByteOrder byteOrder) {
        double[][] doubleMatrix = floatMatrixToDoubleMatrix(matrix);
        return doubleMatrixToByteArray(doubleMatrix, byteOrder);
    }

    /**
     * 二维double数组转short数组（自动限幅）
     *
     * @param matrix    二维double数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return short数组
     */
    public static byte[] doubleMatrixToByteArray(double[][] matrix, ByteOrder byteOrder) {
        ByteBuffer buf = ByteBuffer.allocate(matrix[0].length * 2 * 2);
        buf.order(byteOrder);

        for (int i = 0; i < matrix[0].length; i++) {

            if (matrix[0][i] > 1f) {
                matrix[0][i] = 1f;
            }
            if (matrix[0][i] < -1f) {
                matrix[0][i] = -1f;
            }

            if (matrix[1][i] > 1f) {
                matrix[1][i] = 1f;
            }
            if (matrix[1][i] < -1f) {
                matrix[1][i] = -1f;
            }

            buf.putShort((short) (matrix[0][i] * 32767f));
            buf.putShort((short) (matrix[1][i] * 32767f));

        }
        return buf.array();
    }

    /**
     * byte数组转float数组
     *
     * @param data      byte数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return float数组
     * @throws IllegalArgumentException 如果byte数组长度非法
     */
    public static float[] byteArrayToFloatData(byte[] data, ByteOrder byteOrder) {
        short[] shorts = new short[data.length / 2];
        ByteBuffer.wrap(data).order(byteOrder).asShortBuffer().get(shorts);
        float[] fData = new float[shorts.length];
        for (int i = 0; i < shorts.length; i++) {
            fData[i] = shorts[i] / 32767f;
        }
        return fData;
    }

    /**
     * double数组转float数组
     *
     * @param data double数组
     *
     * @return float数组
     */
    public static float[] doubleDataToFloatData(double[] data) {
        float[] floats = new float[data.length];
        for (int i = 0; i < data.length; i++) {
            double value = data[i];
            floats[i] = (float) value;
        }
        return floats;
    }

    /**
     * 二维float数组转float数组
     *
     * @param matrix 二维float数组
     *
     * @return float数组
     */
    public static float[] floatMatrixToFloatData(float[][] matrix) {
        int numSamples = matrix.length;
        float[] monoArray = new float[numSamples];

        for (int i = 0; i < numSamples; i++) {
            float left = matrix[i][0];
            float right = matrix[i][1];

            // 缩放左右声道，确保它们之和不超过 1.0
            float scale = 1.0f / (Math.abs(left) + Math.abs(right));
            // 避免除以 0
            if (scale > 1.0f) {
                scale = 1.0f;
            }

            monoArray[i] = (left + right) * scale;
        }

        return monoArray;
    }

    /**
     * byte数组转double数组
     *
     * @param data      byte数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return double数组
     */
    public static double[] byteArrayToDoubleData(byte[] data, ByteOrder byteOrder) {
        ByteBuffer buf = ByteBuffer.wrap(data);

        buf.order(byteOrder);
        int i = 0;
        double[] dData = new double[data.length / 2];

        while (buf.remaining() > 0) {
            short s = buf.getShort();
            dData[i] = (double) s / 32767f;
            ++i;
        }
        return dData;
    }

    /**
     * float数组转double数组
     *
     * @param data float数组
     *
     * @return double数组
     */
    public static double[] floatDataToDoubleData(float[] data) {
        double[] doubles = new double[data.length];
        for (int i = 0; i < data.length; i++) {
            float value = data[i];
            doubles[i] = value;
        }
        return doubles;
    }

    /**
     * 二维float数组转double数组
     *
     * @param matrix 二维float数组
     *
     * @return double数组
     */
    public static double[] floatMatrixToDoubleData(float[][] matrix) {
        float[] floats = floatMatrixToFloatData(matrix);
        return floatDataToDoubleData(floats);
    }

    /**
     * byte数组转二维double数组
     *
     * @param data      byte数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return 二维double数组
     */
    public static double[][] byteArrayToDoubleMatrix(byte[] data, ByteOrder byteOrder) {
        ByteBuffer buf = ByteBuffer.wrap(data);

        buf.order(byteOrder);
        int i = 0;
        double[][] dData = new double[2][data.length / 2 / 2];

        while (buf.remaining() > 0) {
            short leftS = buf.getShort();
            short rightS = buf.getShort();

            dData[0][i] = (double) leftS / 32767f;
            dData[1][i] = (double) rightS / 32767f;

            ++i;
        }
        return dData;
    }

    /**
     * 二维double数组转float数组
     *
     * @param matrix 二维double数组
     *
     * @return float数组
     */
    public static double[] doubleMatrixToDoubleData(double[][] matrix) {
        int numSamples = matrix.length;
        double[] monoArray = new double[numSamples];

        for (int i = 0; i < numSamples; i++) {
            double left = matrix[i][0];
            double right = matrix[i][1];

            // 缩放左右声道，确保它们之和不超过 1.0
            double scale = 1.0f / (Math.abs(left) + Math.abs(right));
            // 避免除以 0
            if (scale > 1.0f) {
                scale = 1.0f;
            }

            monoArray[i] = (left + right) * scale;
        }

        return monoArray;
    }

    /**
     * 二维double数组转float数组
     *
     * @param matrix 二维double数组
     *
     * @return float数组
     */
    public static float[] doubleMatrixToFloatData(double[][] matrix) {
        double[] doubles = doubleMatrixToDoubleData(matrix);
        return doubleDataToFloatData(doubles);
    }

    /**
     * byte数组转二维float数组
     *
     * @param data      byte数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     *
     * @return 二维float数组
     * @throws IllegalArgumentException 如果byte数组长度非法
     */
    public static float[][] byteArrayToFloatMatrix(byte[] data, ByteOrder byteOrder) {
        double[][] dData = byteArrayToDoubleMatrix(data, byteOrder);
        return doubleMatrixToFloatMatrix(dData);
    }

    /**
     * 二维float数组转二维double数组
     *
     * @param matrix 二维float数组
     *
     * @return 二维double数组
     */
    public static double[][] floatMatrixToDoubleMatrix(float[][] matrix) {
        double[][] result = new double[matrix.length][matrix[0].length];
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                result[i][j] = matrix[i][j];
            }
        }
        return result;
    }

    /**
     * 二维double数组转二维float数组
     *
     * @param matrix 二维double数组
     *
     * @return 二维float数组
     */
    public static float[][] doubleMatrixToFloatMatrix(double[][] matrix) {
        float[][] result = new float[matrix.length][matrix[0].length];
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                result[i][j] = (float) matrix[i][j];
            }
        }
        return result;
    }

    /**
     * byte数组转short数组
     * @param byteArray byte数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     * @return short数组
     */
    public static short[] byteArrayToShortArray(byte[] byteArray, ByteOrder byteOrder) {
        if (byteArray == null || byteArray.length % 2 != 0) {
            throw new IllegalArgumentException("Byte array must be non-null and have an even length.");
        }

        int shortArrayLength = byteArray.length / 2;
        short[] shortArray = new short[shortArrayLength];

        ByteBuffer buffer = ByteBuffer.wrap(byteArray);
        buffer.order(byteOrder);

        for (int i = 0; i < shortArrayLength; i++) {
            shortArray[i] = buffer.getShort();
        }

        return shortArray;
    }

    /**
     * short数组转byte数组
     * @param shortArray short数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     * @return byte数组
     */
    public static byte[] shortArrayToByteArray(short[] shortArray, ByteOrder byteOrder) {
        if (shortArray == null) {
            throw new IllegalArgumentException("Short array must be non-null.");
        }

        int byteArrayLength = shortArray.length * 2;
        byte[] byteArray = new byte[byteArrayLength];

        ByteBuffer buffer = ByteBuffer.wrap(byteArray);
        buffer.order(byteOrder);

        for (short s : shortArray) {
            buffer.putShort(s);
        }

        return byteArray;
    }

    /**
     * 缩放 short 数组到 byte 数组
     * @param shortArray short 数组
     * @param byteOrder 字节序（ByteOrder.LITTLE_ENDIAN 或 ByteOrder.BIG_ENDIAN）
     * @return byte 数组
     */
    public static byte[] shortArrayToByteArrayScaled(short[] shortArray, ByteOrder byteOrder) {
        if (shortArray == null) {
            throw new IllegalArgumentException("Short array must be non-null.");
        }

        int byteArrayLength = shortArray.length;
        byte[] byteArray = new byte[byteArrayLength];

        for (int i = 0; i < shortArray.length; i++) {
            // 1. 缩放到 byte 的范围 (-128 to 127)
            float normalized = (float) shortArray[i] / 32768.0f;
            byte byteValue = (byte) (normalized * 127.0f);

            byteArray[i] = byteValue;
        }

        return byteArray;
    }

    /**
     * 归一化 [基于峰值 (Peak Normalization)]
     *
     * @param audioData         音频数据
     * @param compressionFactor 压缩因子
     *
     * @return 归一化后的音频数据
     */
    public static float[] normalizeAudioPeak(float[] audioData, float compressionFactor) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        // 1. 找到峰值
        float maxAmplitude = 0.0f;
        for (float sample : audioData) {
            maxAmplitude = Math.max(maxAmplitude, Math.abs(sample));
        }

        // 2. 避免除以零
        if (maxAmplitude == 0.0f) {
            return audioData;
        }

        // 3. 计算增益 (带压缩)
        float gain = 1.0f / maxAmplitude;
        gain = (float) Math.pow(gain, compressionFactor);

        // 4. 应用增益
        float[] normalizedAudio = new float[audioData.length];
        for (int i = 0; i < audioData.length; i++) {
            normalizedAudio[i] = audioData[i] * gain;
        }

        return normalizedAudio;
    }

    /**
     * 归一化 [基于 RMS (Root Mean Square)]
     *
     * @param audioData 音频数据
     * @param targetRms 目标RMS
     *
     * @return 归一化后的音频数据
     */
    public static float[] normalizeAudioRms(float[] audioData, float targetRms) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        // 1. 计算 RMS
        double sumOfSquares = 0.0;
        for (float sample : audioData) {
            sumOfSquares += sample * sample;
        }
        double rms = Math.sqrt(sumOfSquares / audioData.length);

        // 2. 避免除以零
        if (rms == 0.0) {
            return audioData;
        }

        // 3. 计算增益
        float gain = (float) (targetRms / rms);

        // 4. 应用增益
        float[] normalizedAudio = new float[audioData.length];
        for (int i = 0; i < audioData.length; i++) {
            normalizedAudio[i] = audioData[i] * gain;
        }

        return normalizedAudio;
    }

    /**
     * 归一化 [滑动窗口 RMS]
     *
     * @param audioData  音频数据
     * @param windowSize 窗口大小
     * @param targetRms  目标RMS
     *
     * @return 归一化后的音频数据
     */
    public static float[] normalizeAudioRmsWindowed(float[] audioData, int windowSize, float targetRms) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        float[] normalizedAudio = new float[audioData.length];

        for (int i = 0; i < audioData.length; i += windowSize) {
            // 1. 获取窗口数据
            int windowEnd = Math.min(i + windowSize, audioData.length);
            float[] windowData = new float[windowEnd - i];
            System.arraycopy(audioData, i, windowData, 0, windowEnd - i);

            // 2. 计算窗口 RMS
            double sumOfSquares = 0.0;
            for (float sample : windowData) {
                sumOfSquares += sample * sample;
            }
            double rms = Math.sqrt(sumOfSquares / windowData.length);

            // 3. 避免除以零
            float gain = (rms == 0.0) ? 1.0f : (float) (targetRms / rms);

            // 4. 应用增益到窗口
            for (int j = i; j < windowEnd; j++) {
                normalizedAudio[j] = audioData[j] * gain;
            }
        }

        return normalizedAudio;
    }

    /**
     * 归一化 [基于峰值 (Peak Normalization)]
     *
     * @param audioData         音频数据
     * @param compressionFactor 压缩因子
     *
     * @return 归一化后的音频数据
     */
    public static float[][] normalizeAudioPeak(float[][] audioData, float compressionFactor) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        int numChannels = 2;
        int numSamples = audioData.length;

        // 1. 找到左右声道的峰值
        float maxAmplitudeLeft = 0.0f;
        float maxAmplitudeRight = 0.0f;

        for (float[] audioDatum : audioData) {
            maxAmplitudeLeft = Math.max(maxAmplitudeLeft, Math.abs(audioDatum[0]));
            maxAmplitudeRight = Math.max(maxAmplitudeRight, Math.abs(audioDatum[1]));
        }

        // 2. 找到最大的峰值 (可选: 使用左右声道中较大的峰值)
        float maxAmplitude = Math.max(maxAmplitudeLeft, maxAmplitudeRight);

        // 3. 避免除以零
        if (maxAmplitude == 0.0f) {
            return audioData;
        }

        // 4. 计算增益 (带压缩)
        float gain = 1.0f / maxAmplitude;
        gain = (float) Math.pow(gain, compressionFactor);

        // 5. 应用增益到左右声道
        float[][] normalizedAudio = new float[numSamples][numChannels];
        for (int i = 0; i < numSamples; i++) {
            normalizedAudio[i][0] = audioData[i][0] * gain;
            normalizedAudio[i][1] = audioData[i][1] * gain;
        }

        return normalizedAudio;
    }

    /**
     * 归一化 [基于 RMS (Root Mean Square)]
     *
     * @param audioData 音频数据
     * @param targetRms 目标RMS
     *
     * @return 归一化后的音频数据
     */
    public static float[][] normalizeAudioRms(float[][] audioData, float targetRms) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        int numChannels = 2;
        int numSamples = audioData.length;

        // 1. 分别计算左右声道的 RMS
        double sumOfSquaresLeft = 0.0;
        double sumOfSquaresRight = 0.0;

        for (float[] audioDatum : audioData) {
            sumOfSquaresLeft += audioDatum[0] * audioDatum[0];
            sumOfSquaresRight += audioDatum[1] * audioDatum[1];
        }

        double rmsLeft = Math.sqrt(sumOfSquaresLeft / numSamples);
        double rmsRight = Math.sqrt(sumOfSquaresRight / numSamples);

        // 2. 避免除以零
        float gainLeft = (rmsLeft == 0.0) ? 1.0f : (float) (targetRms / rmsLeft);
        float gainRight = (rmsRight == 0.0) ? 1.0f : (float) (targetRms / rmsRight);

        // 3. 应用增益到左右声道
        float[][] normalizedAudio = new float[numSamples][numChannels];
        for (int i = 0; i < numSamples; i++) {
            normalizedAudio[i][0] = audioData[i][0] * gainLeft;
            normalizedAudio[i][1] = audioData[i][1] * gainRight;
        }

        return normalizedAudio;
    }

    /**
     * 归一化 [滑动窗口 RMS]
     *
     * @param audioData  音频数据
     * @param windowSize 窗口大小
     * @param targetRms  目标RMS
     *
     * @return 归一化后的音频数据
     */
    public static float[][] normalizeAudioRmsWindowed(float[][] audioData, int windowSize, float targetRms) {
        if (audioData == null || audioData.length == 0) {
            return audioData;
        }

        int numChannels = 2;
        int numSamples = audioData.length;
        float[][] normalizedAudio = new float[numSamples][numChannels];

        for (int i = 0; i < numSamples; i += windowSize) {
            // 1. 获取窗口数据
            int windowEnd = Math.min(i + windowSize, numSamples);

            // 2. 分别计算左右声道的窗口 RMS
            double sumOfSquaresLeft = 0.0;
            double sumOfSquaresRight = 0.0;

            for (int j = i; j < windowEnd; j++) {
                sumOfSquaresLeft += audioData[j][0] * audioData[j][0];
                sumOfSquaresRight += audioData[j][1] * audioData[j][1];
            }

            int windowLength = windowEnd - i;
            double rmsLeft = Math.sqrt(sumOfSquaresLeft / windowLength);
            double rmsRight = Math.sqrt(sumOfSquaresRight / windowLength);

            // 3. 避免除以零
            float gainLeft = (rmsLeft == 0.0) ? 1.0f : (float) (targetRms / rmsLeft);
            float gainRight = (rmsRight == 0.0) ? 1.0f : (float) (targetRms / rmsRight);

            // 4. 应用增益到左右声道
            for (int j = i; j < windowEnd; j++) {
                normalizedAudio[j][0] = audioData[j][0] * gainLeft;
                normalizedAudio[j][1] = audioData[j][1] * gainRight;
            }
        }

        return normalizedAudio;
    }
}
